package com.yuma.multimodel.controller;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.template.st.StTemplateRenderer;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;
import java.util.Map;

@RestController
public class PromptController {

    @Autowired
    private OllamaChatModel ollamaChatModel;

    @GetMapping("/prompt/model/{name}/{style}")
    public String model(@PathVariable String name, @PathVariable String style) {
        String userText = "请给我讲述三体的内容";
        Message userMessage = new UserMessage(userText);
        String systemText = """
              你是一个热爱科幻小说的的 AI 助手，帮助人们总结科幻小说中心思想。
              你的名字是 {name}。
              你应该用你的名字回复用户的请求，并以一种 {style} 的风格进行回复。
              """;
        SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(systemText);
        Message systemMessage = systemPromptTemplate.createMessage(Map.of("name", name, "style", style));

        Prompt prompt = new Prompt(List.of(userMessage, systemMessage));

        List<Generation> results = ollamaChatModel.call(prompt).getResults();
        return results.get(0).getOutput().getText();
    }

    @GetMapping("/prompt/client/{job}")
    public String client(@PathVariable String job) {
        String answer = ChatClient.create(ollamaChatModel).prompt()
                .user(u -> u
                        .text("告诉我三年 {job} 经验需要掌握的技能.")
                        .param("job", job))
                .call()
                .content();
        return answer;
    }

    // ------------------自定义提示词模板
    @GetMapping("/model/define")
    public String model2() {
        PromptTemplate promptTemplate = PromptTemplate.builder()
                .renderer(StTemplateRenderer.builder().startDelimiterToken('<').endDelimiterToken('>').build())
                .template("推荐我<count>本<type>的书.")
                .build();
        String prompt = promptTemplate.
                render(Map.of("count", 3, "type", "古代军事"));
        return ollamaChatModel.call(prompt);
    }

    @GetMapping("/client/define")
    public String client2() {
        String answer = ChatClient.create(ollamaChatModel).prompt()
                .user(u -> u
                        .text("推荐我<count>本<type>的书.")
                        .param("count", 3)
                        .param("type", "动画"))
                .templateRenderer(StTemplateRenderer.builder().startDelimiterToken('<').endDelimiterToken('>').build())
                .call()
                .content();
        return answer;
    }

    @Value("classpath:prompt.st")
    private Resource resource;

    @GetMapping("/model/file/{type}")
    public String file(@PathVariable String type) {
        SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(resource);
        Message systemMessage = systemPromptTemplate.createMessage(Map.of("type", type));
        Prompt prompt = new Prompt(systemMessage);
        List<Generation> results = ollamaChatModel.call(prompt).getResults();
        return results.get(0).getOutput().getText();
    }

    @GetMapping("/client/file/{type}")
    public String file2(@PathVariable String type) {
        ChatClient  chatClient = ChatClient.builder(ollamaChatModel)
                .defaultSystem(resource)
                .build();
        return chatClient.prompt()
                .system(p -> p.param("type",type))
                .call()
                .content();
    }
}
